package agent

import (
	"sort"
	"strings"
	"sync"
	"time"
)

// Memory 记忆管理器
type Memory struct {
	shortTerm    []Message              // 短期记忆（当前会话）
	longTerm     map[string][]MemoryItem // 长期记忆（按会话ID分组）
	maxSize      int                    // 最大记忆容量
	mutex        sync.RWMutex           // 读写锁
}

// NewMemory 创建新的记忆管理器
func NewMemory(maxSize int) *Memory {
	return &Memory{
		shortTerm: make([]Message, 0),
		longTerm:  make(map[string][]MemoryItem),
		maxSize:   maxSize,
	}
}

// AddMessage 添加消息到短期记忆
func (m *Memory) AddMessage(message Message) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.shortTerm = append(m.shortTerm, message)

	// 如果超过最大容量，移除最旧的消息
	if len(m.shortTerm) > m.maxSize {
		m.shortTerm = m.shortTerm[1:]
	}

	// 同时添加到长期记忆
	m.addToLongTerm(message)
}

// addToLongTerm 添加到长期记忆
func (m *Memory) addToLongTerm(message Message) {
	item := MemoryItem{
		ID:        generateID(),
		Content:   message.Content,
		Type:      "conversation",
		Metadata: map[string]string{
			"role": message.Role,
		},
		Timestamp: message.Timestamp,
		SessionID: message.SessionID,
		Score:     m.calculateImportanceScore(message),
	}

	if _, exists := m.longTerm[message.SessionID]; !exists {
		m.longTerm[message.SessionID] = make([]MemoryItem, 0)
	}

	m.longTerm[message.SessionID] = append(m.longTerm[message.SessionID], item)

	// 保持长期记忆的大小限制
	if len(m.longTerm[message.SessionID]) > m.maxSize*2 {
		// 按重要性排序，保留重要的记忆
		sort.Slice(m.longTerm[message.SessionID], func(i, j int) bool {
			return m.longTerm[message.SessionID][i].Score > m.longTerm[message.SessionID][j].Score
		})
		m.longTerm[message.SessionID] = m.longTerm[message.SessionID][:m.maxSize]
	}
}

// calculateImportanceScore 计算重要性评分
func (m *Memory) calculateImportanceScore(message Message) float64 {
	score := 0.5 // 基础分数

	// 用户消息通常比助手消息重要
	if message.Role == "user" {
		score += 0.2
	}

	// 包含特定关键词的消息更重要
	importantKeywords := []string{"重要", "记住", "提醒", "不要忘记", "关键", "必须"}
	content := strings.ToLower(message.Content)
	for _, keyword := range importantKeywords {
		if strings.Contains(content, keyword) {
			score += 0.3
			break
		}
	}

	// 较长的消息可能包含更多信息
	if len(message.Content) > 100 {
		score += 0.1
	}

	// 时间衰减：较新的消息稍微重要一些
	hoursSinceCreation := time.Since(message.Timestamp).Hours()
	if hoursSinceCreation < 24 {
		score += 0.1
	}

	return score
}

// GetRecentMessages 获取最近的消息
func (m *Memory) GetRecentMessages(count int) []Message {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if count >= len(m.shortTerm) {
		return m.shortTerm
	}

	return m.shortTerm[len(m.shortTerm)-count:]
}

// GetSessionMessages 获取指定会话的所有消息
func (m *Memory) GetSessionMessages(sessionID string) []Message {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	var messages []Message
	for _, msg := range m.shortTerm {
		if msg.SessionID == sessionID {
			messages = append(messages, msg)
		}
	}

	return messages
}

// SearchMemory 搜索记忆
func (m *Memory) SearchMemory(query string, sessionID string, limit int) []MemoryItem {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	query = strings.ToLower(query)
	var results []MemoryItem

	// 搜索指定会话的长期记忆
	if items, exists := m.longTerm[sessionID]; exists {
		for _, item := range items {
			if m.matchesQuery(item, query) {
				results = append(results, item)
			}
		}
	}

	// 按相关性和重要性排序
	sort.Slice(results, func(i, j int) bool {
		scoreI := m.calculateRelevanceScore(results[i], query)
		scoreJ := m.calculateRelevanceScore(results[j], query)
		return scoreI > scoreJ
	})

	// 限制结果数量
	if limit > 0 && len(results) > limit {
		results = results[:limit]
	}

	return results
}

// matchesQuery 检查记忆项是否匹配查询
func (m *Memory) matchesQuery(item MemoryItem, query string) bool {
	content := strings.ToLower(item.Content)
	return strings.Contains(content, query)
}

// calculateRelevanceScore 计算相关性评分
func (m *Memory) calculateRelevanceScore(item MemoryItem, query string) float64 {
	content := strings.ToLower(item.Content)
	query = strings.ToLower(query)

	score := item.Score // 基于重要性评分

	// 精确匹配加分
	if strings.Contains(content, query) {
		score += 0.5
	}

	// 关键词匹配
	queryWords := strings.Fields(query)
	contentWords := strings.Fields(content)
	
	matchCount := 0
	for _, qWord := range queryWords {
		for _, cWord := range contentWords {
			if strings.Contains(cWord, qWord) {
				matchCount++
				break
			}
		}
	}

	if len(queryWords) > 0 {
		score += float64(matchCount) / float64(len(queryWords)) * 0.3
	}

	// 时间因素：较新的记忆稍微重要一些
	hoursSinceCreation := time.Since(item.Timestamp).Hours()
	if hoursSinceCreation < 24 {
		score += 0.1
	} else if hoursSinceCreation < 168 { // 一周内
		score += 0.05
	}

	return score
}

// AddFact 添加事实性记忆
func (m *Memory) AddFact(sessionID, content string, metadata map[string]string) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	item := MemoryItem{
		ID:        generateID(),
		Content:   content,
		Type:      "fact",
		Metadata:  metadata,
		Timestamp: time.Now(),
		SessionID: sessionID,
		Score:     0.8, // 事实性记忆通常比较重要
	}

	if _, exists := m.longTerm[sessionID]; !exists {
		m.longTerm[sessionID] = make([]MemoryItem, 0)
	}

	m.longTerm[sessionID] = append(m.longTerm[sessionID], item)
}

// AddPreference 添加用户偏好记忆
func (m *Memory) AddPreference(sessionID, content string, metadata map[string]string) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	item := MemoryItem{
		ID:        generateID(),
		Content:   content,
		Type:      "preference",
		Metadata:  metadata,
		Timestamp: time.Now(),
		SessionID: sessionID,
		Score:     0.9, // 用户偏好非常重要
	}

	if _, exists := m.longTerm[sessionID]; !exists {
		m.longTerm[sessionID] = make([]MemoryItem, 0)
	}

	m.longTerm[sessionID] = append(m.longTerm[sessionID], item)
}

// GetMemoryStats 获取记忆统计信息
func (m *Memory) GetMemoryStats() map[string]interface{} {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	totalLongTerm := 0
	for _, items := range m.longTerm {
		totalLongTerm += len(items)
	}

	return map[string]interface{}{
		"short_term_count": len(m.shortTerm),
		"long_term_count":  totalLongTerm,
		"sessions_count":   len(m.longTerm),
		"max_size":         m.maxSize,
	}
}

// ClearSession 清除指定会话的记忆
func (m *Memory) ClearSession(sessionID string) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	// 清除短期记忆中的会话消息
	var filteredShortTerm []Message
	for _, msg := range m.shortTerm {
		if msg.SessionID != sessionID {
			filteredShortTerm = append(filteredShortTerm, msg)
		}
	}
	m.shortTerm = filteredShortTerm

	// 清除长期记忆
	delete(m.longTerm, sessionID)
}

// generateID 生成唯一ID
func generateID() string {
	return time.Now().Format("20060102150405") + "_" + randomString(6)
}

// randomString 生成随机字符串
func randomString(length int) string {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, length)
	for i := range b {
		b[i] = charset[time.Now().UnixNano()%int64(len(charset))]
	}
	return string(b)
}
